from CalTqFixOffset_Gauss import BinData, PolyTwoGauss
import math
import pickle
import numpy
from scipy import interpolate
from scipy.optimize import minimize
from scipy.optimize import curve_fit
from matplotlib import pyplot as PLT

# NEEDS A TON OF WORK


def PolyTwoGauss(p, x):           
        # offset = distance between peaks
        # p[0], p[1], p[2] = amplitude , center_x, sigma: for peak A
        # p[3], p[4] = amplitude, sigma: for peak B
        # p[5] = polynomial offset
        # p[6], p[7], p[8] = ax^2 + bx + c
    offset = 52
    y = numpy.zeros(len(x[0]))
    for i, j in enumerate(x[0]):
        y[i] = p[0]*math.exp(-0.5*((j-p[1])/p[2])**2.) + p[3]*math.exp(-0.5*((j-(p[1]+offset))/p[4])**2.) + p[6]*(j - p[5])**2 + p[7]*(j - p[5]) + p[8]
        if y[i] <= 1e-5: 
            y[i] = 1e-5
    return y

def BinData(x, y):
    bins = 200.0
    max_d = 400.0
    xdata = numpy.linspace(0, max_d, bins)
    ydata = numpy.zeros(bins)

    for j, k in enumerate(xdata):
        pos_torque = []
        x0 = k - (max_d / bins) / 2.
        x1 = k + (max_d / bins) / 2.
        for i, p in enumerate(x):
           if p > x0 and p <= x1:
              pos_torque.append(y[i])
        if len(pos_torque) > 0:
            ydata[j] += max(pos_torque)
    
    return xdata, ydata

def PlotFit_mod(run, outdir = 'fitplots/'):
    """ Plots torque vs position data and Gaussian fit to that data """
    date = run.Datetime().strftime('%Y%m%d')
    direction = run.Direction()
    source = run.Source()
    new_parms = run.NewParms()
    position = run.NewPosition()
    torque = run.Torque()
    station = run.Station()
    
    new_pos = []
    if position[0] != 0: 
        p0 = position[0]
        for j in position:
            new_pos.append(j - p0)
    else:
        new_pos = position
            
    xdata, ydata = BinData(new_pos, torque)
    
    if station == 'S2':
        xbin = xdata[140:170]
        ybin = ydata[140:170]
    else:
        xbin = xdata[140:-10]
        ybin = ydata[140:-10]
    
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax1 = fig.add_axes(box)
    ax1.set_ylabel('Torque (in.oz)')
    ax1.set_xlabel('Position (cm)')
    maxy = max(torque)
    maxx = max(position)
    miny = 0
    
    ## Add station locations
    if maxx >= 328.9:
        ax1.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
    if maxx >= 382.6:
        ax1.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
    if maxx >= 435.7:
        ax1.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
    if maxx >= 491.7:
        ax1.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
    if maxx >= 595.8:
        ax1.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
        ax1.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
    
    ## Add bend locations
    if maxx >= 47.8:
        ax1.axvspan(47.8, 93.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bellows', xy = (48, miny), color = 'g')
    if maxx >= 136.4:
        ax1.axvspan(136.4, 144.6, color = 'g', alpha = 0.2)
        ax1.annotate('First outer bend', xy = (136, miny), color = 'g')
    if maxx >= 196.5:
        ax1.axvspan(196.5, 218.0, color = 'g', alpha = 0.2)
        ax1.annotate('The "S"', xy = (212, miny), color = 'g')
    if maxx >= 228.3:
        ax1.axvspan(228.3, 241.5, color = 'g', alpha = 0.2)
    if maxx >= 274.6:
        ax1.axvspan(274.6, 284.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A5', xy = (274, miny), color = 'g')
    if maxx >= 309.0:
        ax1.axvspan(309.0, 318.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A6', xy = (309, miny), color = 'g')
    if maxx >= 358.8:
        ax1.axvspan(358.8, 370.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A7', xy = (359, miny), color = 'g')
    if maxx >= 414.5:
        ax1.axvspan(414.5, 426.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A8', xy = (414, miny), color = 'g')
    if maxx >= 445.3:
        ax1.axvspan(445.3, 455.3, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A9', xy = (445, miny), color = 'g')
    if maxx >= 467.5:
        ax1.axvspan(467.5, 479.5, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A10', xy = (467, miny), color = 'g')
    if maxx >= 523.8:
        ax1.axvspan(523.8, 535.8, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A11', xy = (523, miny), color = 'g')
    if maxx >= 571.6:
        ax1.axvspan(571.6, 583.6, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A12', xy = (571, miny), color = 'g')
    if maxx >= 627.9:
        ax1.axvspan(627.9, 639.9, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A13', xy = (627, miny), color = 'g')
    if maxx >= 652.7:
        ax1.axvspan(652.7, 660.7, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A14', xy = (652, miny), color = 'g')
    if maxx >= 678.5:
        ax1.axvspan(678.5, 681.1, color = 'g', alpha = 0.2)
        ax1.annotate('Bend A15', xy = (678, miny), color = 'g')
    
    PLT.title(date + ' Source ' + source + ' ' + direction + ' Offset Fix by Model Method')
    ax1.grid()
    ax1.plot(xdata, ydata, '-', color = 'maroon', linewidth = 1, label = 'Binned Original Data')
    ax1.plot(xbin, PolyTwoGauss(new_parms, [xbin]), '-', color = 'g', linewidth = 1, label = 'Poly-Two-Gauss')
    #ax1.plot(run.NewPosition(), torque, '-', color = 'blue', linewidth = 1, label = 'Corrected Offset')
    h, l = ax1.get_legend_handles_labels()
    PLT.legend(h, l, 'upper left')
    fig.savefig(outdir + date + '_' + direction + '_' + source + '_' + station + '_modelfit' + '.png')
    PLT.close(fig)
    